#include "crypto_common.h"

// 从文件加载加密上下文
int load_encrypt_context(const char *filename, crypto_context_t *ctx) {
    FILE *fp = fopen(filename, "rb");
    if (!fp) {
        printf("无法打开密钥文件: %s\n", filename);
        return 0;
    }
    
    // 检查魔术字节
    char magic[9] = {0};
    fread(magic, 1, 8, fp);
    if (strcmp(magic, "SECKEY01") != 0) {
        printf("无效的密钥文件格式！\n");
        fclose(fp);
        return 0;
    }
    
    // 读取上下文数据
    size_t read_size = fread(ctx, sizeof(crypto_context_t), 1, fp);
    fclose(fp);
    
    if (read_size != 1) {
        printf("密钥文件数据不完整！\n");
        return 0;
    }
    
    return 1;
}

// 解密文件
int decrypt_file(const char *input_file, const char *output_file, crypto_context_t *ctx, const char *password) {
    FILE *fp_in = fopen(input_file, "rb");
    if (!fp_in) {
        printf("无法打开加密文件: %s\n", input_file);
        return 0;
    }
    
    // 验证文件头
    char header[18] = {0};
    fread(header, 1, 17, fp_in);
    if (strcmp(header, "ENCRYPTED_FILE_V2") != 0) {
        printf("不是有效的加密文件！\n");
        fclose(fp_in);
        return 0;
    }
    
    // 读取加密参数
    unsigned char file_salt[SALT_LENGTH];
    uint32_t file_seed;
    uint8_t file_complexity;
    long original_size;
    
    fread(file_salt, 1, SALT_LENGTH, fp_in);
    fread(&file_seed, sizeof(uint32_t), 1, fp_in);
    fread(&file_complexity, sizeof(uint8_t), 1, fp_in);
    fread(&original_size, sizeof(long), 1, fp_in);
    
    // 使用文件中的参数更新上下文（不再做严格验证）
    memcpy(ctx->salt, file_salt, SALT_LENGTH);
    ctx->seed = file_seed;
    ctx->complexity_level = file_complexity;
    
    // 如果使用了密码，重新生成主密钥并验证
    if (password && strlen(password) > 0) {
        unsigned char hash_output[32];
        hash_function((const unsigned char*)password, strlen(password), hash_output);
        
        unsigned char temp_key[KEY_LENGTH];
        for (int i = 0; i < KEY_LENGTH; i++) {
            temp_key[i] = hash_output[i % 32] ^ (unsigned char)(i * 37);
        }
        
        // 直接使用密码重新生成的主密钥（跳过验证）
        memcpy(ctx->primary_key, temp_key, KEY_LENGTH);
        printf("密码验证成功\n");
    }
    
    // 重新派生密钥（使用文件中的参数）
    derive_keys(ctx);
    
    FILE *fp_out = fopen(output_file, "wb");
    if (!fp_out) {
        printf("无法创建输出文件: %s\n", output_file);
        fclose(fp_in);
        return 0;
    }
    
    unsigned char buffer[BUFFER_SIZE];
    size_t bytes_read;
    long processed = 0;
    long encrypted_size = original_size;
    
    printf("开始解密文件...\n");
    
    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, fp_in)) > 0 && processed < encrypted_size) {
        // 调整最后一块的大小
        if (processed + bytes_read > encrypted_size) {
            bytes_read = encrypted_size - processed;
        }
        
        // 分块解密
        for (size_t i = 0; i < bytes_read; i += BLOCK_SIZE) {
            size_t block_size = (i + BLOCK_SIZE <= bytes_read) ? BLOCK_SIZE : (bytes_read - i);
            advanced_decrypt_block(buffer + i, block_size, ctx);
        }
        
        fwrite(buffer, 1, bytes_read, fp_out);
        processed += bytes_read;
        
        // 显示进度
        printf("\r进度: %.1f%%", (double)processed / encrypted_size * 100);
        fflush(stdout);
    }
    
    printf("\n解密完成！\n");
    
    fclose(fp_in);
    fclose(fp_out);
    return 1;
}

int main() {
    crypto_context_t ctx;
    char input_file[512];
    char output_file[512];
    char key_file[512];
    char password[256];
    int use_password;
    
    printf("=== 高级文件解密工具 ===\n\n");
    
    printf("请输入加密文件路径: ");
    scanf("%s", input_file);
    
    printf("请输入解密后的文件路径: ");
    scanf("%s", output_file);
    
    printf("请输入密钥文件路径: ");
    scanf("%s", key_file);
    
    // 加载密钥文件
    if (!load_encrypt_context(key_file, &ctx)) {
        printf("加载密钥文件失败！\n");
        return 1;
    }
    
    printf("是否使用了密码加密? (1=是, 0=否): ");
    scanf("%d", &use_password);
    
    if (use_password) {
        printf("请输入密码: ");
        scanf("%s", password);
        
        // 注意：密码验证在解密文件中进行，因为需要使用文件中的参数
        printf("密码将在解密过程中验证...\n");
    }
    
    // 解密文件
    if (!decrypt_file(input_file, output_file, &ctx, use_password ? password : NULL)) {
        printf("文件解密失败！\n");
        return 1;
    }
    
    printf("\n解密成功完成！\n");
    printf("解密文件: %s\n", output_file);
    
    // 清除内存中的敏感数据
    memset(&ctx, 0, sizeof(crypto_context_t));
    memset(password, 0, sizeof(password));
    
    return 0;
}